/**
 * Copyright (c) iwindplus Technologies Co., Ltd.2024-2030, All rights reserved.
 */

package com.iwindplus.gateway.server.filter;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.net.Ipv4Util;
import com.iwindplus.base.domain.constant.CommonConstant.HeaderConstant;
import com.iwindplus.base.domain.constant.CommonConstant.SymbolConstant;
import com.iwindplus.base.domain.dto.MessageBaseDTO;
import com.iwindplus.base.domain.enums.BaseEnum;
import com.iwindplus.base.domain.enums.BizCodeEnum;
import com.iwindplus.base.domain.enums.OperateTypeEnum;
import com.iwindplus.base.domain.exception.BizException;
import com.iwindplus.gateway.server.client.MgtClient;
import com.iwindplus.gateway.server.domain.constant.GatewayConstant.FilterConstant;
import com.iwindplus.gateway.server.domain.property.GatewayProperty;
import com.iwindplus.gateway.server.util.GatewayUtil;
import com.iwindplus.mgt.domain.dto.system.IpBlackListChangeDTO;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.ApplicationArguments;
import org.springframework.boot.ApplicationRunner;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

/**
 * IP黑名单过滤器.
 *
 * @author zengdegui
 * @since 2024-9-24
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class IpBlackListFilter implements Ordered, GlobalFilter, ApplicationRunner {

    private final GatewayProperty property;
    private final MgtClient mgtClient;
    private final AtomicReference<Set<String>> snapshot = new AtomicReference<>(Set.of());

    @Override
    public int getOrder() {
        return FilterConstant.FILTER_IP_BLACK_LIST_ORDER;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 开关未启用跳过当前过滤器
        if (Boolean.FALSE.equals(this.property.getIpBlackList().getEnabled())) {
            return chain.filter(exchange);
        }

        final long start = System.nanoTime();
        final String realIp = exchange.getRequest().getHeaders().getFirst(HeaderConstant.REAL_IP);

        return Mono.justOrEmpty(snapshot.get())
            .filter(set -> !set.isEmpty())
            .switchIfEmpty(loadIpBlackList())
            .flatMap(list -> checkThenChain(exchange, chain, list, realIp, start));
    }

    private Mono<Void> checkThenChain(ServerWebExchange exchange, GatewayFilterChain chain, Set<String> ipList, String ip, long start) {
        if (this.checkIp(ipList, ip)) {
            log.warn("IP [{}] 命中黑名单，拦截路径: {}", ip, exchange.getRequest().getURI().getPath());
            return GatewayUtil.asyncPublishErrorLog(exchange, property.getLog(),
                    BizCodeEnum.IP_IS_BLACKLISTED.getBizStatus())
                .then(Mono.error(new BizException(BizCodeEnum.IP_IS_BLACKLISTED, new Object[]{ip})));
        }

        return chain.filter(exchange)
            .doFinally(signal -> log.info("{} execute cost={} ms", this.getClass().getSimpleName(),
                TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - start)));
    }

    private Mono<Set<String>> loadIpBlackList() {
        return mgtClient.listIp()
            .defaultIfEmpty(List.of())
            .map(this::getStringList)
            .doOnNext(snapshot::set)
            .doOnNext(list -> log.info("IP黑名单加载完成，条目数={}", list.size()))
            .subscribeOn(Schedulers.boundedElastic());
    }

    private Set<String> getStringList(List<String> list) {
        return new HashSet<>(list);
    }

    /**
     * 刷新黑名单.
     *
     * @param message 消息
     */
    public void refreshBlackList(MessageBaseDTO<IpBlackListChangeDTO> message) {
        if (Objects.isNull(message) || Objects.isNull(message.getOperateType()) || Objects.isNull(message.getData())) {
            return;
        }

        final OperateTypeEnum operateType = BaseEnum.fromValue(message.getOperateType(), OperateTypeEnum.class);
        final IpBlackListChangeDTO data = message.getData();
        Set<String> oldSet = snapshot.get();
        Set<String> newSet = new HashSet<>(oldSet);
        switch (operateType) {
            case ADD:
                if (CollUtil.isNotEmpty(data.getNewIp())) {
                    newSet.addAll(data.getNewIp());
                }
                break;
            case DELETE:
                if (CollUtil.isNotEmpty(data.getOldIp())) {
                    newSet.removeAll(data.getOldIp());
                }
                break;
            case MODIFY:
                if (CollUtil.isNotEmpty(data.getOldIp()) && CollUtil.isNotEmpty(data.getNewIp())) {
                    newSet.removeAll(data.getOldIp());
                    newSet.addAll(data.getNewIp());
                }
                break;
            default:
                throw new BizException(BizCodeEnum.UNSUPPORTED_TYPE, new Object[]{message.getOperateType()});
        }
        snapshot.set(Set.copyOf(newSet));
        log.info("IP黑名单刷新完成，条目数={}", newSet.size());
    }

    private boolean checkIp(Set<String> ipList, String ip) {
        Set<String> direct = new HashSet<>(16);
        Set<String> cidr = new HashSet<>(16);

        for (String ipStr : ipList) {
            (ipStr.contains(SymbolConstant.SLASH) ? cidr : direct).add(ipStr);
        }

        if (direct.contains(ip)) {
            return true;
        }
        if (!cidr.isEmpty()) {
            return cidr.stream().anyMatch(c -> Ipv4Util.matches(c, ip));
        }
        return false;
    }

    @Override
    public void run(ApplicationArguments args) throws Exception {
        this.loadIpBlackList().subscribe();
    }

}
